Skip to content

Support torch.distributed.scatter collective #9365

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 8 commits into from
Jun 24, 2025
Merged

Support torch.distributed.scatter collective #9365

merged 8 commits into from
Jun 24, 2025

Conversation

bfolie
Copy link
Collaborator

@bfolie bfolie commented Jun 16, 2025 •
edited
Loading

#9315

XLA doesn't have a distributed Scatter op but we can put dummy tensor lists on the non-source rank and use reduce_scatter

@bfolie bfolie requested review from bhavya01 and pgmoka June 16, 2025 17:02
bfolie
bfolie commented Jun 17, 2025
View reviewed changes
@bfolie bfolie mentioned this pull request Jun 17, 2025
Open
pgmoka
pgmoka reviewed Jun 17, 2025
View reviewed changes
Copy link
Collaborator

@pgmoka pgmoka left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mostly questions, and a requests for extra documentation

pgmoka
pgmoka reviewed Jun 17, 2025
View reviewed changes
@bfolie bfolie requested review from pgmoka and ghpvnist June 18, 2025 17:32
ghpvnist
ghpvnist approved these changes Jun 20, 2025
View reviewed changes
pgmoka
pgmoka approved these changes Jun 20, 2025
View reviewed changes
Copy link
Collaborator

@pgmoka pgmoka left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up seems good. Let me know if you have any questions on https://github.com/pytorch/xla/pull/9365/files#r2151351304.

Otherwise, LGTM

pgmoka
pgmoka reviewed Jun 20, 2025
View reviewed changes
Copy link
Collaborator

@pgmoka pgmoka left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up seems good. Let me know if you have any questions on https://github.com/pytorch/xla/pull/9365/files#r2151351304.

Otherwise, LGTM.

One minor thing: I believe the tests failing are due to flakyness. Can you confirm?

@bfolie
Copy link
Collaborator Author

bfolie commented Jun 24, 2025

One minor thing: I believe the tests failing are due to flakyness. Can you confirm?

The TPU test failure is a known flake and was cleared up by re-running.

The Torchprime e2e test failure is probably real but because PRs from forks aren't exercising that test it's very difficult to tell where the failure comes from

@bfolie bfolie merged commit 01db65d into master Jun 24, 2025
23 of 24 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@pgmoka pgmoka pgmoka approved these changes

@ghpvnist ghpvnist ghpvnist approved these changes

@bhavya01 bhavya01 Awaiting requested review from bhavya01

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@bfolie @pgmoka @ghpvnist